import os
import sys
import cv2
import json
import shutil
import numpy as np
import argparse
"""
代码功能说明：
create_output_folders 函数：
创建输出文件夹结构，包括图像和掩码文件夹。
save_output_images 函数：
遍历输入路径下的所有数据集。
从 JSON 文件中加载注释信息。
如果是训练集，直接复制图像到输出路径。
如果是测试集，调整图像大小并保存。
save_output_masks 函数：
遍历输入路径下的所有数据集。
从 JSON 文件中加载注释信息。
创建掩码图像，根据注释中的多边形点填充掩码。
保存掩码图像。
main 函数：
获取命令行参数。
创建输出文件夹。
遍历训练集和测试集，保存图像和掩码。
get_args 函数：
使用 argparse 解析命令行参数。
"""

# 创建输出文件夹
def create_output_folders(out_path, parts):
    # 遍历每个部分（训练集或测试集）
    for part in parts:
        # 创建图像文件夹
        if not os.path.exists(os.path.join(out_path, part, 'images')):
            os.makedirs(os.path.join(out_path, part, 'images'))
        # 创建掩码文件夹
        if not os.path.exists(os.path.join(out_path, part, 'masks')):
            os.makedirs(os.path.join(out_path, part, 'masks'))


# 保存输出图像
def save_output_images(in_path, out_path, part):
    # 获取输入路径下的所有数据集
    datasets = os.listdir(in_path)
    for dataset in datasets:
        dataset_path = os.path.join(in_path, dataset)
        # 找到 JSON 文件
        json_path = [f for f in os.listdir(dataset_path) if f.endswith('.json')]
        # 加载注释文件
        annotations = json.load(open(os.path.join(dataset_path, json_path[0])))

        for a in annotations:
            if annotations[a]["regions"]:
                if part == 'train':
                    # 获取图像名称
                    image_name = annotations[a]["filename"]
                    # 将图像复制到输出路径
                    shutil.copy2(os.path.join(dataset_path, image_name), os.path.join(out_path, "images", dataset + "_" + image_name))
                else:
                    # 获取图像名称
                    image_name = annotations[a]["filename"]
                    # 读取图像并调整大小
                    image = cv2.imread(os.path.join(in_path, dataset, image_name))
                    image = cv2.resize(image, (720, 1280))
                    # 保存调整大小后的图像
                    cv2.imwrite(os.path.join(out_path, "images", dataset + "_" + image_name), image)


# 保存输出掩码
def save_output_masks(in_path, out_path, part):
    # 获取输入路径下的所有数据集
    datasets = os.listdir(in_path)
    for dataset in datasets:
        dataset_path = os.path.join(in_path, dataset)
        # 找到 JSON 文件
        json_path = [f for f in os.listdir(dataset_path) if f.endswith('.json')]
        # 加载注释文件
        annotations = json.load(open(os.path.join(dataset_path, json_path[0])))

        for a in annotations:
            if annotations[a]["regions"]:
                # 获取图像名称
                image_name = annotations[a]["filename"]
                # 定义掩码路径
                mask_path = os.path.join(out_path, "masks", dataset + "_" + image_name)

                # 创建掩码文件
                mask = np.zeros((1280, 720), np.uint8)
                count = 1
                for region in annotations[a]["regions"]:
                    print(dataset)
                    # 检查区域是否为苹果
                    if region['region_attributes']['class'].lower() == 'Apple'.lower():
                        # 获取多边形的点
                        x = region["shape_attributes"]["all_points_x"]
                        y = region["shape_attributes"]["all_points_y"]

                        pts = np.array(np.column_stack((x, y)), np.int32)
                        pts = pts.reshape((-1, 1, 2))
                        # 在掩码上填充多边形
                        cv2.fillPoly(mask, [pts], count)
                        count += 1

                # 保存掩码
                cv2.imwrite(mask_path, mask)


# 主函数
def main():
    args = get_args()
    base_path = args.base_path  # 基础路径
    out_path = args.out_path  # 输出路径
    parts = ['train', 'test']  # 数据集部分

    # 创建输出文件夹
    create_output_folders(out_path, parts)

    for part in parts:
        read_path = os.path.join(base_path, part)  # 读取路径
        write_path = os.path.join(out_path, part)  # 写入路径

        # 保存输出图像和掩码
        save_output_images(read_path, write_path, part)
        save_output_masks(read_path, write_path, part)


# 获取命令行参数
def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--base_path', type=str, help='基础路径')
    parser.add_argument('--out_path', type=str, help='输出路径')
    return parser.parse_args()


if __name__ == "__main__":
    main()